{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "header-cell",
   "metadata": {},
   "source": [
    "# 🎨 Qwen 图像编辑器 - Gradio版\n",
    "\n",
    "基于 Qwen-Image-Edit 模型的交互式图像编辑应用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install-deps",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 安装必要的依赖\n",
    "!pip install gradio modelscope diffusers transformers accelerate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "import torch\n",
    "from PIL import Image\n",
    "import os\n",
    "from diffusers import QwenImageEditPipeline\n",
    "from modelscope import snapshot_download\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup-device",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置设备和数据类型\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "    torch_dtype = torch.bfloat16\n",
    "    print(f\"✅ 使用 GPU: {torch.cuda.get_device_name()}\")\n",
    "    print(f\"💾 GPU 内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB\")\n",
    "else:\n",
    "    device = \"cpu\"\n",
    "    torch_dtype = torch.float32\n",
    "    print(\"⚠️ 使用 CPU（建议使用GPU以获得更好性能）\")\n",
    "\n",
    "print(f\"🔧 设备: {device}, 数据类型: {torch_dtype}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "download-model",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下载并加载 Qwen-Image-Edit 模型\n",
    "model_id = \"Qwen/Qwen-Image-Edit\"\n",
    "local_dir = './models/Qwen-Image-Edit'\n",
    "\n",
    "# 检查模型是否已存在\n",
    "if not os.path.exists(local_dir):\n",
    "    print(f\"📥 开始下载 {model_id} 模型...\")\n",
    "    os.makedirs(os.path.dirname(local_dir), exist_ok=True)\n",
    "    snapshot_download(model_id, local_dir=local_dir)\n",
    "    print(f\"✅ 模型下载完成: {local_dir}\")\n",
    "else:\n",
    "    print(f\"✅ 模型已存在: {local_dir}\")\n",
    "\n",
    "# 加载管道\n",
    "print(\"🔄 正在加载图像编辑管道...\")\n",
    "pipeline = QwenImageEditPipeline.from_pretrained(\n",
    "    local_dir, \n",
    "    torch_dtype=torch_dtype,\n",
    "    use_safetensors=True,\n",
    "    device_map=\"balanced\"\n",
    ")\n",
    "print(pipeline)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edit-function",
   "metadata": {},
   "outputs": [],
   "source": [
    "def edit_image_with_prompt(image, prompt, seed=42, steps=50, cfg_scale=4.0):\n",
    "    \"\"\"\n",
    "    使用提示词编辑图像\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # 处理输入图像\n",
    "        if isinstance(image, np.ndarray):\n",
    "            image = Image.fromarray(image)\n",
    "        \n",
    "        if image.mode != \"RGB\":\n",
    "            image = image.convert(\"RGB\")\n",
    "        \n",
    "        # 编辑参数\n",
    "        inputs = {\n",
    "            \"image\": image,\n",
    "            \"prompt\": prompt,\n",
    "            \"generator\": torch.manual_seed(seed),\n",
    "            \"true_cfg_scale\": cfg_scale,\n",
    "            \"negative_prompt\": \" \",\n",
    "            \"num_inference_steps\": steps,\n",
    "        }\n",
    "        \n",
    "        # 执行编辑\n",
    "        with torch.inference_mode():\n",
    "            output = pipeline(**inputs)\n",
    "            edited_image = output.images[0]\n",
    "        \n",
    "        return edited_image\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ 编辑图像时出错: {str(e)}\")\n",
    "        return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "gradio-interface",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建 Gradio 界面\n",
    "def create_image_editor():\n",
    "    with gr.Blocks(title=\"Qwen 图像编辑器\", theme=gr.themes.Soft()) as demo:\n",
    "        gr.Markdown(\n",
    "            \"\"\"\n",
    "            # 🎨 Qwen 智能图像编辑器\n",
    "            \n",
    "            使用自然语言描述您想要的编辑效果，AI 将为您智能编辑图像！\n",
    "            \n",
    "            **使用步骤：** 上传图像 → 输入编辑提示词 → 点击编辑按钮\n",
    "            \"\"\"\n",
    "        )\n",
    "        \n",
    "        with gr.Row():\n",
    "            with gr.Column(scale=1):\n",
    "                # 输入区域\n",
    "                input_image = gr.Image(\n",
    "                    label=\"📷 上传图像\",\n",
    "                    type=\"pil\",\n",
    "                    height=350\n",
    "                )\n",
    "                \n",
    "                prompt_input = gr.Textbox(\n",
    "                    label=\"✏️ 编辑提示词\",\n",
    "                    placeholder=\"请详细描述您想要的编辑效果，例如：将背景改为蓝天白云，添加一只可爱的小猫...\",\n",
    "                    lines=3\n",
    "                )\n",
    "                \n",
    "                # 快速示例按钮\n",
    "                gr.Markdown(\"**🚀 快速示例：**\")\n",
    "                with gr.Row():\n",
    "                    btn1 = gr.Button(\"🎨 动漫风格\", size=\"sm\")\n",
    "                    btn2 = gr.Button(\"🌅 日落背景\", size=\"sm\")\n",
    "                    btn3 = gr.Button(\"🐱 添加小猫\", size=\"sm\")\n",
    "                    btn4 = gr.Button(\"❄️ 冬季雪景\", size=\"sm\")\n",
    "                \n",
    "                # 参数设置\n",
    "                with gr.Accordion(\"🎛️ 高级参数\", open=False):\n",
    "                    seed_input = gr.Slider(0, 1000, value=42, step=1, label=\"随机种子\")\n",
    "                    steps_input = gr.Slider(10, 100, value=50, step=5, label=\"推理步数\")\n",
    "                    cfg_scale_input = gr.Slider(1.0, 10.0, value=4.0, step=0.5, label=\"CFG缩放\")\n",
    "                \n",
    "                edit_button = gr.Button(\"🎨 开始编辑\", variant=\"primary\", size=\"lg\")\n",
    "            \n",
    "            with gr.Column(scale=1):\n",
    "                # 输出区域\n",
    "                output_image = gr.Image(label=\"✨ 编辑结果\", height=350)\n",
    "        \n",
    "        # 示例提示词说明\n",
    "        with gr.Accordion(\"💡 提示词示例\", open=False):\n",
    "            gr.Markdown(\n",
    "                \"\"\"\n",
    "                **风格转换：**\n",
    "                - \"Transform into anime/manga style with vibrant colors\"\n",
    "                - \"Convert to oil painting style\"\n",
    "                - \"将这张照片转换为水彩画风格\"\n",
    "                \n",
    "                **场景编辑：**\n",
    "                - \"Change background to beautiful sunset\"\n",
    "                - \"Add snow and winter atmosphere\"\n",
    "                - \"将白天改为夜晚，添加灯光\"\n",
    "                \n",
    "                **对象编辑：**\n",
    "                - \"Add a cute cat in the foreground\"\n",
    "                - \"Remove the person from image\"\n",
    "                - \"在图像中添加一朵美丽的花\"\n",
    "                \"\"\"\n",
    "            )\n",
    "        \n",
    "        # 编辑处理函数\n",
    "        def process_edit(image, prompt, seed, steps, cfg_scale, progress=gr.Progress(track_tqdm=True)):\n",
    "            if image is None:\n",
    "                return None\n",
    "            if not prompt.strip():\n",
    "                return None\n",
    "            \n",
    "            try:\n",
    "                # 执行编辑\n",
    "                edited_image = edit_image_with_prompt(image, prompt, seed, steps, cfg_scale)\n",
    "                return edited_image\n",
    "                \n",
    "            except Exception as e:\n",
    "                return None\n",
    "        \n",
    "        \n",
    "        # 绑定事件\n",
    "        edit_button.click(\n",
    "            fn=process_edit,\n",
    "            inputs=[input_image, prompt_input, seed_input, steps_input, cfg_scale_input],\n",
    "            outputs=[output_image]\n",
    "        )\n",
    "        \n",
    "        # 示例按钮绑定\n",
    "        examples = [\n",
    "            \"Transform this image into anime/manga style with vibrant colors and soft lighting\",\n",
    "            \"Change the background to a beautiful sunset with warm golden and orange colors\",\n",
    "            \"Add a cute cat sitting naturally in the foreground of this image\",\n",
    "            \"Transform this scene to winter with snow covering everything\"\n",
    "        ]\n",
    "        \n",
    "        for btn, example in zip([btn1, btn2, btn3, btn4], examples):\n",
    "            btn.click(fn=lambda x=example: x, outputs=[prompt_input])\n",
    "    \n",
    "    return demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "launch-app",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"🚀 启动 Qwen 图像编辑器...\")\n",
    "    \n",
    "demo = create_image_editor()\n",
    "\n",
    "# 启动 Gradio 应用\n",
    "demo.launch(\n",
    "    share=False,          # 创建公共链接，便于远程访问\n",
    "    server_name=\"0.0.0.0\",  # 允许外部访问\n",
    "    server_port=6006,    # 端口号\n",
    "    show_error=True,     # 显示详细错误信息\n",
    "    debug=True          # 生产环境建议关闭debug\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
